Async grpo OpenEnv harness rollout by AmineDiro · Pull Request #6420 · huggingface/trl

AmineDiro · 2026-07-16T13:39:34Z

WIP

local sandbox: https://gist.github.com/AmineDiro/efc1b8f5b654690311f6f7ae6df7ad67
opencode harness: https://gist.github.com/AmineDiro/4c236418255c8cc6c697eacddf886981

from __future__ import annotations

import argparse

from datasets import Dataset
from transformers import AutoTokenizer

from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer
from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker

from opencode_harness import MODEL, build_dataset, build_factory


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default=MODEL)
    p.add_argument("--vllm-url", default="http://localhost:8000")
    p.add_argument("--num-generations", type=int, default=4)
    p.add_argument("--max-inflight", type=int, default=4)  # per-rollout uuid sandbox + free proxy port -> concurrent
    p.add_argument("--max-completion-length", type=int, default=16384)
    p.add_argument("--max-steps", type=int, default=10)
    p.add_argument("--n-prompts", type=int, default=32)
    p.add_argument("--project", default="async-grpo-opencode-mbpp")
    p.add_argument("--output-dir", default="async_grpo_opencode_mbpp")
    args = p.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    dataset = Dataset.from_list(build_dataset(n_prompts=args.n_prompts, seed=0))

    worker = HarnessRolloutWorker(
        harness_session_factory=build_factory(),
        harness_adapter=None,  # loop-owning: opencode runs its own loop; TRL reads the proxy trace
        model_name=args.model,
        dataset=dataset,
        reward_funcs=[],
        processing_class=tokenizer,
        num_generations=args.num_generations,
        max_inflight_tasks=args.max_inflight,
        vllm_server_url=args.vllm_url,
        max_tokens=args.max_completion_length,
        temperature=1.0,
        fork_threshold_tokens=1024,
        log_completions=True,
        num_completions_to_print=2,
    )

    config = AsyncGRPOConfig(
        output_dir=args.output_dir,
        save_strategy="no",
        per_device_train_batch_size=4,
        num_generations=args.num_generations,
        gradient_accumulation_steps=1,
        max_completion_length=args.max_completion_length,
        max_steps=args.max_steps,
        learning_rate=1e-5,
        vllm_server_base_url=args.vllm_url,
        report_to="trackio",
        project=args.project,
        log_completions=True,
    )

    trainer = AsyncGRPOTrainer(
        model=args.model,
        args=config,
        train_dataset=dataset,
        processing_class=tokenizer,
        rollout_worker=worker,
    )
    trainer.train()


if __name__ == "__main__":
    main()

Note

Medium Risk
Touches the experimental async GRPO rollout/scoring path (new reward column and return signatures) and runs untrusted agent code via local subprocess sandboxes; failures are mostly isolated but concurrent rollouts increase operational complexity.

Overview
Adds OpenEnv harness support to Async GRPO so rollouts can be driven by external agents that own their tool loop, instead of TRL’s built-in vLLM turn loop.

A new HarnessRolloutWorker / _HarnessRolloutLoop runs OpenEnv sessions (white-box via HarnessAdapter or loop-owning via proxy trace), rebuilds TurnRecords from captured token ids/logprobs, verifies in the sandbox, and scores via an optional rollout_reward_fn. Hooks train_turn_fn and agent_turn_fn filter which trace turns become training rows. Failed sessions degrade to unscorable rollouts instead of killing the worker.

Async rollout plumbing is extended: RolloutGroup.rollout_rewards, _generate_one returns a sixth rollout_reward slot, _provides_rollout_reward lets loops score without reward_funcs, the child process accepts a pluggable _loop_cls, and RolloutWorkerProtocol allows multiprocessing.Queue buffers.

The examples/scripts/openenv/opencode.py script demonstrates end-to-end training: LocalSubprocessSandboxBackend (path remap from /home/user, process-group cleanup, template hardlink clones), DeepCoder held-out stdin/stdout verifier, per-session free proxy ports, and opencode-specific reward/turn filtering wired into AsyncGRPOTrainer.

^{Reviewed by Cursor Bugbot for commit 9015404. Bugbot is set up for automated code reviews on this repo. Configure here.}

cursor · 2026-07-16T14:42:32Z

+            reward = float(verify.env_reward) if verify.env_reward is not None else None
+            sequences = _chain_to_sequences(turns, rollout_id, self._fork_threshold_tokens)
+            completion_ids = [tid for turn in turns for tid in turn.output_ids]
+            return completion, completion_ids, sequences, tool_call_count, tool_failure_count, reward


Empty traces skew advantages

High Severity

When the proxy trace is empty or _trace_output_ids recovers no token ids, _chain_to_sequences yields no training rows, but session.verify still supplies a rollout_reward that enters the group's advantage normalization. Sibling generations are then trained against rewards from conversations that contribute no samples.

Additional Locations (1)

trl/experimental/async_grpo/openenv_harness.py#L144-L156

^{Reviewed by Cursor Bugbot for commit 284da95. Configure here.}

When the proxy trace is empty or _trace_output_ids recovers no token ids

no idea in what case this is possible

I dont think thats real case

bot-ci-comment · 2026-07-16T15:33:44Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

qgallouedec · 2026-07-20T04:22:28Z

Thanks I think it looks good overall. I didn't try myself, but from reading the edits, nothing much to say.

Check the cursor review, it's likely the reason of the CI failure. Once the CI green I approve and we merge

sergiopaniego

thanks!!
approving but with 2 ideas:

we could add a short harness-rollout section to openenv.md in this PR: white-box vs loop-owning, --return-tokens-as-token-ids, reward via session.verify() → harness_reward.
we should add an example script. I think this could be added once we add support for HF sandbox + harnesses. wdyt? it's in the roadmap here -> huggingface/OpenEnv#940

sergiopaniego · 2026-07-22T09:27:02Z

Ran this branch loop-owning with OpenCode on HF Jobs; two things surfaced (both line up with the Cursor comments above, now with a concrete repro):

1. _messages_from_trace crashes on empty choices. The last proxy-trace entry can have choices == [] (an aborted/blank final request from the agent); choices[0] then raises IndexError, which kills the rollout child and stops training:

File ".../trl/experimental/async_grpo/openenv_harness.py", in _messages_from_trace
    content = last["response"]["choices"][0]["message"].get("content", "")
IndexError: list index out of range

Minimal guard:

choices = last["response"].get("choices") or []
content = choices[0]["message"].get("content", "") if choices else ""

(plus rebuilding the final assistant tool_calls, per the line-182 comment). With this guard applied, the same run trains to completion.

2. Rollout child doesn't exit cleanly. Confirmed the [RANK 0] Child did not exit within 15s; terminating path (the session-pool shutdown() already flagged). Side effect with a remote sandbox backend: in-flight sandboxes then leak, because session.close() never runs when the child is force-terminated.

sergiopaniego · 2026-07-22T14:53:01Z

Heads-up on another rough edge hit while training a loop-owning harness (Pi coding agent) on this branch, separate from the empty-choices one:

print_prompt_completions_sample → format_entry (trl/trainer/utils.py) does t.append(msg["content"]), which assumes content is a str. Agents that emit structured content (a list of parts, e.g. [{"type": "text", "text": "..."}] — Pi does this) trip rich's Text.append with TypeError: Only str or Text can be appended to Text, crashing the rollout worker when log_completions=True.

Repro-side workaround: coerce structured content to text before append, e.g.

c = msg["content"]
t.append(c if isinstance(c, str) else "".join(
    p.get("text", "") if isinstance(p, dict) else str(p) for p in c))

Happy to open a small PR if useful.

qgallouedec · 2026-07-22T15:50:18Z

@sergiopaniego can you check #5309, would it solve the issue? I can revamp it you think it makes sense

AmineDiro · 2026-07-22T16:15:06Z

print_prompt_completions_sample → format_entry (trl/trainer/utils.py) does t.append(msg["content"]), which assumes content is a str. Agents that emit structured content (a list of parts, e.g. [{"type": "text", "text": "..."}] — Pi does this) trip rich's Text.append with TypeError: Only str or Text can be appended to Text, crashing the rollout worker when log_completions=True.

@sergiopaniego thanks for catching this, but print_prompt_completions_sample might not be useful in harness training. I've opened this PR for tracking traces in trackio directly 👍🏼

…ssion cleanup

cursor · 2026-07-22T16:32:29Z

+        # it should retry with backoff (ideally inside OpenEnv's factory). For now a failed create drops the rollout as
+        # unscorable, which shrinks the effective group size - a high create-failure rate silently weakens the signal.
+        try:
+            session = self._factory.create(prompt, seed=seed, episode_id=rollout_id)


Group seed never varies

High Severity

_HarnessRolloutLoop._generate_one takes group_id and uses it as the session seed, but the parent generate loop never passes group_id, so every rollout keeps the default 0. Seed-driven factories therefore hand every group the same task for the whole run.

Additional Locations (1)

trl/experimental/async_grpo/async_rollout_worker.py#L555-L556

^{Reviewed by Cursor Bugbot for commit 06b6083. Configure here.}

cursor · 2026-07-22T16:32:29Z

+    last = entries[-1]
+    choices = (last.get("response") or {}).get("choices") or []
+    content = choices[0]["message"].get("content", "") if choices else ""
+    return list(last["request"].get("messages") or []) + [{"role": "assistant", "content": content}]


Structured content crashes logging

High Severity

Loop-owning completions keep agent message content as-is, including structured lists. With log_completions=True, print_prompt_completions_sample appends that value to rich Text and raises TypeError, which kills the score loop and stops the rollout worker.

^{Reviewed by Cursor Bugbot for commit 06b6083. Configure here.}

AmineDiro · 2026-07-22T16:38:09Z

In this commit : 06b6083

fixed @sergiopaniego you can rerun if possible :

Sandbox leaks fixed: Added live-session tracking and explicit pool teardown to ensure remote sandboxes are cleanly killed on shutdown.
Tool counts fixed: Removed double-counting bugs; there is now a single source of truth for tool metrics.

OpenEnv Workarounds @burtenshaw @adithya-s-k @sergiopaniego

I had to work around some OpenEnv fundamental type /API limitation. I noted them as TODO(@openenv) in this PR:

Missing Types & Methods: Declared local types (TraceEntry, loop-owning protocols, generic factories) to patch gaps in OpenEnv's API.
Heuristic Call Filtering: Manually filtering out background LLM calls (like summarization) because OpenEnv doesn't tag trace purposes.

So what I just pushed is

The harness is now completely framework-agnostic. Opencode-specific logic is moved to the app layer via three injectable hooks: rollout_reward_fn, train_turn_fn, and agent_turn_fn.
- rollout_reward_fn(HarnessRolloutOutcome) -> float | None — reward shaping (default: env reward, unshaped).
- train_turn_fn(HarnessTurn) -> bool — which turns to reinforce ie to actually train on.
- agent_turn_fn(list[TraceEntry]) -> list[TraceEntry] — which trace entries are real agent turns
Strict Typing: Replaced raw dicts with strict, documented data models (e.g., HarnessRolloutOutcome, HarnessTurn).

I also added this : cf #6349 @qgallouedec

Deterministic Rollouts: Tasks are now seeded by group_id (not the prompt), ensuring all generations within a group get the exact same task.

qgallouedec

lgtm, just a bug (not sure) on group id

qgallouedec · 2026-07-22T16:48:46Z

+        # In-flight sessions, so `_run_loops` can close them on stop (see there). set ops are atomic under the GIL.
+        self._live_sessions: set = set()
+
+    async def _generate_one(self, prompt, tool_dict, tools, group_id=0):


group_id isn't yet wired, right?

qgallouedec · 2026-07-22T16:51:48Z

                        )
                        pending_completed[group_id] = 0

                    task = asyncio.create_task(self._generate_one(prompt, tool_dict=tool_dict, tools=tools))


see https://github.com/huggingface/trl/pull/6420/changes#r3632113776, shouldn't we have group_id here?

qgallouedec · 2026-07-22T16:54:24Z

+def _tool_failure_count(entries: list[TraceEntry]) -> int:
+    """Best-effort tool-failure count across the real agent turns (`entries`). The agent (not TRL) executed the tools,
+    so we only see their free-text RESULTS, which come back as `role="tool"` messages in a later request; a failure can
+    only be inferred from error-looking result text. The last agent request holds the fullest message list; dedupe by
+    (name, content) so a result echoed across requests is not counted twice."""
+    seen, failures = set(), 0
+    for msg in (entries[-1]["request"] if entries else {}).get("messages") or []:
+        if msg.get("role") != "tool":
+            continue
+        key = (msg.get("name"), str(msg.get("content"))[:200])
+        if key in seen:
+            continue
+        seen.add(key)
+        text = str(msg.get("content") or "").lower()
+        if any(w in text for w in ("error", "failed", "traceback", "exception")):
+            failures += 1
+    return failures


quite fragile, consider dropping it.
if you want to keep it, let's keep in mind that we should re-visit later when openenv exposes real tool-result status

qgallouedec · 2026-07-22T16:56:59Z

    policy version from the shared `mp.Value` (`model_version_value`).
    """

+    # Class attribute  declares that this loop produces its own per-rollout reward


Suggested change

# Class attribute declares that this loop produces its own per-rollout reward

# Class attribute: declares that this loop produces its own per-rollout reward

cursor

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 4 total unresolved issues (including 3 from previous reviews).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 9015404. Configure here.}

cursor · 2026-07-22T17:05:58Z

+        return []
+    last = entries[-1]
+    choices = (last.get("response") or {}).get("choices") or []
+    content = choices[0]["message"].get("content", "") if choices else ""


Fragile final-choice message access

Medium Severity

Empty choices is handled, but a non-empty choice missing message still does choices[0]["message"] and raises. That happens after turns are already built, so the outer handler discards an otherwise trainable rollout as unscorable.

Additional Locations (1)

trl/experimental/async_grpo/openenv_harness.py#L216-L239

^{Reviewed by Cursor Bugbot for commit 9015404. Configure here.}

AmineDiro requested a review from qgallouedec July 16, 2026 13:39

AmineDiro added 2 commits July 16, 2026 14:04

impl openenv harness rollout worker

b458bcc

clean up openenv harness rollout worker comments

1903378

AmineDiro force-pushed the async-grpo-harness-rollout branch from dabfebb to 1903378 Compare July 16, 2026 14:07

simplify child loop-class swap via _loop_cls attribute

284da95

AmineDiro marked this pull request as ready for review July 16, 2026 14:38

AmineDiro requested review from adithya-s-k, burtenshaw and sergiopaniego July 16, 2026 14:40

cursor Bot reviewed Jul 16, 2026

View reviewed changes

Merge branch 'main' into async-grpo-harness-rollout

4a4af23

cursor Bot reviewed Jul 16, 2026

View reviewed changes

Comment thread tests/experimental/test_async_grpo_trainer.py

AmineDiro changed the title ~~Async grpo harness rollout~~ Async grpo OpenEnv harness rollout Jul 16, 2026

AmineDiro and others added 2 commits July 20, 2026 08:40

Fix async_grpo unit tests for 6-tuple rollout API

f91730c

Merge branch 'main' into async-grpo-harness-rollout

e6d759d

sergiopaniego approved these changes Jul 20, 2026

View reviewed changes

This was referenced Jul 20, 2026

RFC 006: Agentic RL through Harness Interception — token-faithful traces for TRL huggingface/OpenEnv#941

Open

Add HFSandboxBackend for the OpenCode harness huggingface/OpenEnv#998

Open

sergiopaniego mentioned this pull request Jul 22, 2026

Add pi_env: Pi coding-agent environment (HF sandbox + transparent-proxy logprobs) huggingface/OpenEnv#999

Open

Add harness rollout worker policy hooks, typed trace, and graceful se…

06b6083

…ssion cleanup

cursor Bot reviewed Jul 22, 2026

View reviewed changes

qgallouedec approved these changes Jul 22, 2026

View reviewed changes

Add opencode loop-owning AsyncGRPO example (local sandbox)

9015404

cursor Bot reviewed Jul 22, 2026

View reviewed changes

	# Class attribute declares that this loop produces its own per-rollout reward
	# Class attribute: declares that this loop produces its own per-rollout reward

Uh oh!

Conversation

AmineDiro commented Jul 16, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot Jul 16, 2026

Choose a reason for hiding this comment

Empty traces skew advantages

Uh oh!

qgallouedec Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

AmineDiro Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

bot-ci-comment Bot commented Jul 16, 2026

Uh oh!

Uh oh!

qgallouedec commented Jul 20, 2026

Uh oh!

sergiopaniego left a comment

Choose a reason for hiding this comment

Uh oh!

sergiopaniego commented Jul 22, 2026

Uh oh!

sergiopaniego commented Jul 22, 2026

Uh oh!

qgallouedec commented Jul 22, 2026

Uh oh!

AmineDiro commented Jul 22, 2026

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Group seed never varies

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Structured content crashes logging

Uh oh!

AmineDiro commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

OpenEnv Workarounds @burtenshaw @adithya-s-k @sergiopaniego

Uh oh!

qgallouedec left a comment

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

qgallouedec Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor Bot Jul 22, 2026

Choose a reason for hiding this comment

Fragile final-choice message access

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AmineDiro commented Jul 16, 2026 •

edited by cursor Bot

Loading

AmineDiro commented Jul 22, 2026 •

edited

Loading